#include <kalman.h>

using namespace KalmanJoint;
using namespace std;

void MultiSensorKalmanFilter::update(const SensorNumber& sensor_num) {
    setSensorFilterState(sensor_num);

    state_estimate_ = A_ * state_estimate_  + B_ * input_ + Q_ * w_;
    P_ = A_ * P_ * A_.transpose() + Q_;

    Eigen::MatrixXd& H = sensor_info_[sensor_num].sensor_H;
    Eigen::MatrixXd& R = sensor_info_[sensor_num].sensor_R;
    Eigen::MatrixXd S = H * P_ * H.transpose() + R;
    Eigen::MatrixXd K = P_ * H.transpose();
    K = S.ldlt().solve(K);// 直接求解线性方程，防止S不可逆以及使用伪逆降低效率引入误差
    state_estimate_ += K * (getSensorData(sensor_num) - H * state_estimate_);
    P_ = (Eigen::MatrixXd::Identity(state_size_, state_size_) - K * H) * P_;
}

void MultiSensorKalmanFilter::updateAllSensor(){
    if (getUsableFlag(IMU)) update(IMU);
    setTimeStep(0.);      // 不迭代模型，仅更新其他观测数据源
    if (getUsableFlag(RGBD_Apriltag)) update(RGBD_Apriltag);
    if (getUsableFlag(RGBD_YOLO_RFRNN)) update(RGBD_YOLO_RFRNN);
}

void MultiSensorKalmanFilter::updateAllData(SensorNumber base_sensor, const Eigen::VectorXd& base_sensor_data, const long double curr_time){
    setSensorData(base_sensor, base_sensor_data);
    setSensorTimeStamp(base_sensor, curr_time);
    setUsableFlag(base_sensor, true);
    setUpdateFlag(base_sensor, true);
    for (const auto& sensor : sensor_info_) {
        const SensorNumber sensor_num = sensor.first;
        if (sensor_num != base_sensor){
            if(getSensorTimeStamp(sensor_num) < getNewestTimeStamp(sensor_num)){
                setSensorData(sensor_num, getNewestData(sensor_num));
                setSensorTimeStamp(sensor_num, getNewestTimeStamp(sensor_num));
                setUpdateFlag(sensor_num, true);
            }
        }
    }
}

void MultiSensorKalmanFilter::updateAkf(const SensorNumber& sensor_num, const Eigen::VectorXd& measurement){
    state_estimate_ = A_ * state_estimate_ + B_ * input_ + Q_ * w_;
    P_ = A_ * P_ * A_.transpose() + Q_;

    Eigen::MatrixXd& H = sensor_info_[sensor_num].sensor_H;
    Eigen::MatrixXd& R = sensor_info_[sensor_num].sensor_R;
    Eigen::Matrix<double, 3, 1> r = sensor_info_[sensor_num].sensor_residual;
    if(!r.isZero()){
        sensor_info_[sensor_num].sensor_R = alpha_ * R + (1-alpha_) * (r *r.transpose() +  H * P_ * H.transpose());
    }
    Eigen::MatrixXd S = H * P_ * H.transpose() + R;
    Eigen::Matrix<double, 3, 1> y = measurement - H * state_estimate_;
    Eigen::MatrixXd K = P_ * H.transpose();
    K = S.ldlt().solve(K);
    Q_ = alpha_ * Q_ + (1-alpha_) * (K * y * y.transpose() * K.transpose());
    state_estimate_ += K * y;
    sensor_info_[sensor_num].sensor_residual = y;
    P_ = (Eigen::MatrixXd::Identity(state_size_, state_size_) - K * H) * P_;
}

void MultiSensorKalmanFilter::addSensor(const SensorParams& params) {
    if (sensor_info_.find(params.sensor_num) != sensor_info_.end()) {
        throw std::runtime_error(
                "Sensor number " + std::to_string(params.sensor_num) + " already exists in sensor_info_.");
        return;
    }

    SensorParams sensor_params = params;
    sensor_params.sensor_residual = Eigen::Matrix<double, 3, 1>::Zero();
    sensor_params.measure_date = Eigen::VectorXd::Zero(params.sensor_H.rows());
    sensor_params.newest_measure_date = Eigen::VectorXd::Zero(params.sensor_H.rows());
    sensor_params.time_stamp = 0.;
    sensor_params.newest_time_stamp = 0.;
    sensor_params.update_flag = false;
    sensor_params.usable_flag = false;
    sensor_params.sensor_B = params.sensor_H * this->getTimeStep();

    int rows = params.sensor_H.rows();
    sensor_params.sensor_w = Eigen::VectorXd::Zero(rows);
    for (int i = 0; i < rows; ++i) {
        if (i < params.sensor_H.cols() && params.sensor_H(i, i) == 1.0) {
            sensor_params.sensor_w(i) = 1.0;
        }
    }
    sensor_info_[params.sensor_num] = sensor_params;
}

void MultiSensorKalmanFilter::addHisData(const long double curr_time) {
    std::unordered_map<SensorNumber, std::pair<bool, Eigen::VectorXd>> sensor_angle;
    for (const auto& sensor : sensor_info_) {
        sensor_angle[sensor.first] =  std::make_pair(getUsableFlag(sensor.first), getSensorData(sensor.first));
    }
    HistoryData his_data = {state_estimate_, input_, curr_time, sensor_angle};

    if (history_data_.size() >= QUEUE_LENGTH) {
        history_data_.erase(history_data_.begin()); // 移除第一个元素
    }

    // 添加新的历史数据到末尾
    history_data_.push_back(his_data);
}

void MultiSensorKalmanFilter::updateFromHis(const SensorNumber base_sensor) {
    const long double base_timestamp = getSensorTimeStamp(base_sensor);
    long double last_time_ms = base_timestamp;

    // 找到最接近 base_timestamp 的历史数据
    auto closest_imu_it = history_data_.begin();
    for (auto it = history_data_.begin(); it != history_data_.end(); ++it) {
        if (it->timestamp >= base_timestamp) {
            break;
        }
        closest_imu_it = it;
    }

    // 如果没有找到匹配的历史数据，直接返回
    if (closest_imu_it == history_data_.end()) {
        return;
    }

    // 使用最接近的历史数据初始化状态
    setStateEstimate(closest_imu_it->angle);
    setInput(closest_imu_it->velocity);
    setTimeStep(closest_imu_it->timestamp, base_timestamp);
    update(base_sensor);

    // 从找到的历史数据开始更新传感器
    for (auto it = closest_imu_it+1; it != history_data_.end(); ++it) {
        setInput(it->velocity);
        setTimeStep(last_time_ms, it->timestamp);

        for (const auto& sensor : sensor_info_) {
            const SensorNumber sensor_num = sensor.first;
            if (sensor_num != base_sensor) {
                if (it->sensor_angle[sensor_num].first) {
                    setSensorData(sensor_num, it->sensor_angle[sensor_num].second);
                }
            }
        }

        updateAllSensor();
        it->angle = getStateEstimate();
        it->sensor_angle[base_sensor].first = true;
        it->sensor_angle[base_sensor].second = getSensorData(base_sensor);
        last_time_ms = it->timestamp;
    }
}